#ifndef STX_ALGORITHM_PARTITION_HPP
#define STX_ALGORITHM_PARTITION_HPP

#include <stx/algorithm/common.hpp>
#include <stx/algorithm/swap.hpp>

namespace stx {

template <class Predicate, class ForwardIterator>
inline ForwardIterator
partition_impl(
    ForwardIterator first,
    ForwardIterator last,
    Predicate pred,
    std::forward_iterator_tag)
{
    while (1) {
        if (first == last) {
            return first;
        }
        if (!pred(*first)) {
            break;
        }
        ++first;
    }
    for (ForwardIterator p = first; ++p != last;) {
        if (pred(*p)) {
            stx::swap(*first, *p);
            ++first;
        }
    }
    return first;
}

template <class Predicate, class BidirectionalIterator>
inline BidirectionalIterator
partition_impl(
    BidirectionalIterator first,
    BidirectionalIterator last,
    Predicate pred,
    std::bidirectional_iterator_tag)
{
    while (1) {
        while (1) {
            if (first == last) {
                return first;
            }
            if (!pred(*first)) {
                break;
            }
            ++first;
        }
        do {
            if (first == --last) {
                return first;
            }
        } while (!pred(*last));
        stx::swap(*first, *last);
        ++first;
    }
}

template <class ForwardIterator, class Predicate>
inline ForwardIterator
partition(ForwardIterator first, ForwardIterator last, Predicate pred)
{
    return stx::partition_impl<typename add_lvalue_reference<Predicate>::type>(
        first,
        last,
        pred,
        typename std::iterator_traits<ForwardIterator>::iterator_category());
}

} // namespace stx

#endif // STX_ALGORITHM_PARTITION_HPP
